from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.exceptions import RequestValidationError
from datetime import datetime
import asyncio
import atexit

# 导入配置和工具
from config import config
from exceptions import ExceptionHandler, BaseCustomException, ErrorCode
from utils import LogManager

# 导入数据库
from database import create_tables

# 导入进程管理器
from process_manager import process_manager

# 导入边侧管理器
from edge.edge_manager import edge_manager

# 导入所有路由模块
from routers import health, training, models, datasets, nodes, samples, system, config as config_router, scheduling, test_sets

# 初始化配置和日志
config.ensure_directories()
LogManager.initialize_logging(config.log.level)
main_logger = LogManager.get_logger("main")

# 创建FastAPI应用
app = FastAPI(
    title=config.api.title,
    description=config.api.description,
    version=config.api.version,
    docs_url="/docs",
    redoc_url="/redoc"
)

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=config.api.cors_origins,
    allow_credentials=True,
    allow_methods=config.api.cors_methods,
    allow_headers=config.api.cors_headers,
)

# 创建数据库表
create_tables()

# 注册路由
app.include_router(health.router)
app.include_router(training.router)
app.include_router(models.router)
app.include_router(datasets.router)
app.include_router(nodes.router)
app.include_router(samples.router)
app.include_router(system.router)
app.include_router(config_router.router)
app.include_router(scheduling.router)
app.include_router(test_sets.router)

# 初始化云端节点数据
@app.on_event("startup")
async def startup_event():
    """应用启动时的初始化操作"""
    main_logger.info("应用启动，开始初始化云端节点数据...")
    try:
        # 导入服务
        from services import NodeService, DatasetService, ModelService
        from services.scheduling_service import SchedulingService
        from database import get_db
        
        # 获取数据库会话
        db_gen = get_db()
        db = next(db_gen)
        
        # 初始化服务
        node_service = NodeService()
        dataset_service = DatasetService()
        model_service = ModelService()
        scheduling_service = SchedulingService(db)
        
        # 从配置文件初始化节点数据
        node_service.initialize_nodes_from_config(db)
        main_logger.info("云端节点数据初始化完成")
        
        # 从datasets文件夹加载数据集配置
        loaded_count = dataset_service.load_profiles_from_datasets(db)
        main_logger.info(f"数据集配置初始化完成，共加载 {loaded_count} 个配置")
        
        # 初始化预训练模型
        model_count = model_service.initialize_pretrained_models(db)
        main_logger.info(f"预训练模型初始化完成，共加载 {model_count} 个模型")
        
        # 初始化调度策略
        strategy_count = scheduling_service.initialize_strategies_from_config(db)
        main_logger.info(f"调度策略初始化完成，共加载 {strategy_count} 个策略")
        
        # 初始化进程管理器（确保已完全初始化）
        main_logger.info("进程管理器初始化完成")
        
        # 注册清理函数
        atexit.register(cleanup_on_exit)
        
    except Exception as e:
        main_logger.error(f"初始化失败: {str(e)}")
    finally:
        db.close()

@app.on_event("shutdown")
async def shutdown_event():
    """应用关闭时的清理操作"""
    main_logger.info("应用关闭，开始清理资源...")
    try:
        # 清理EdgeManager中的所有任务
        edge_manager.cleanup_all_tasks()
        
        # 清理进程管理器中的所有资源
        process_manager.cleanup_all()
        
        main_logger.info("资源清理完成")
    except Exception as e:
        main_logger.error(f"清理资源时出错: {str(e)}")

def cleanup_on_exit():
    """退出时的清理函数"""
    try:
        main_logger.info("程序退出，执行最终清理...")
        
        # 创建新的事件循环来处理异步清理
        try:
            loop = asyncio.get_event_loop()
            if loop.is_running():
                # 如果循环正在运行，创建任务
                loop.create_task(async_cleanup())
            else:
                # 如果循环已停止，运行清理
                loop.run_until_complete(async_cleanup())
        except RuntimeError:
            # 如果没有事件循环，创建新的
            asyncio.run(async_cleanup())
            
    except Exception as e:
        print(f"最终清理时出错: {str(e)}")

async def async_cleanup():
    """异步清理函数"""
    try:
        edge_manager.cleanup_all_tasks()
        process_manager.cleanup_all()
    except Exception as e:
        print(f"异步清理时出错: {str(e)}")

# 导入统一响应格式工具
from response_utils import error_response

# 全局异常处理器
@app.exception_handler(BaseCustomException)
async def custom_exception_handler(request: Request, exc: BaseCustomException):
    """处理自定义异常"""
    main_logger.error(f"自定义异常: {exc.error_code.name} - {exc.message}")
    return JSONResponse(
        status_code=ExceptionHandler._get_http_status_code(exc.error_code),
        content=error_response(
            message=exc.message,
            error_code=str(exc.error_code.value),
            error_type=exc.error_code.name,
            details=exc.details
        )
    )

@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
    """处理请求验证异常"""
    # 安全地序列化错误信息
    safe_errors = []
    for error in exc.errors():
        safe_error = {
            "type": error.get("type", "unknown"),
            "loc": error.get("loc", []),
            "msg": str(error.get("msg", "")),
            "input": str(error.get("input", ""))
        }
        if "url" in error:
            safe_error["url"] = error["url"]
        safe_errors.append(safe_error)
    
    main_logger.error(f"请求验证失败: {safe_errors}")
    return JSONResponse(
        status_code=400,
        content=error_response(
            message="请求参数验证失败",
            error_code=str(ErrorCode.VALIDATION_ERROR.value),
            error_type="VALIDATION_ERROR",
            details={"validation_errors": safe_errors}
        )
    )

@app.exception_handler(Exception)
async def general_exception_handler(request: Request, exc: Exception):
    """处理通用异常"""
    main_logger.error(f"未处理异常: {type(exc).__name__} - {str(exc)}", exc_info=True)
    return JSONResponse(
        status_code=500,
        content=error_response(
            message="系统内部错误",
            error_code=str(ErrorCode.INTERNAL_ERROR.value),
            error_type="INTERNAL_ERROR",
            details={"error": str(exc)}
        )
    )

if __name__ == "__main__":
    import uvicorn
    main_logger.info("启动边缘云端AI训练系统")
    uvicorn.run(
        app, 
        host=config.api.host, 
        port=config.api.port,
        log_level=config.log.level.lower()
    )
